from dna.Parameters import *
from dna.PolymorphismTables import PolymorphismTables
import dna.SiteFrequencySpectrum as sfs
import numpy.matlib as np
from math import sqrt
import scipy.stats.mstats as mstats
import dna.infiniteSitesTools as ist
import matplotlib.pyplot as plt
from dna.PolyspectOperator import *
from dna.PowerTools import Power

#test
NUM_OF_SIMULATIONS = 100
NUM_OF_CHROMOSOMES = 6

POWER_CALCULATION_SIZE = 10

N = 10000
MU = 1e-8
LENGTH = 25000
THETA = 4.0 * N * MU * LENGTH
S = 50

p = Population(NUM_OF_CHROMOSOMES, N, LENGTH, MU, 0.0, S)
d = Demography(N)

simParams = SimulationParameters( p, d )
print simParams

# Get starting point (Tajima's D)
myTD = ist.getOp_TajimasD( NUM_OF_CHROMOSOMES )
print myTD

# Initialize and normalize the Polyspect
polyOp = Polyspect(myTD, simParams, NUM_OF_SIMULATIONS )

#ssList = polyOp.operateOnAll( polyOp.mySFSs )

print polyOp.getRejectionValues(0.05)
myPower = Power(polyOp)

# Now to iterate through demography
altDemography = Demography(N)
newSize = N * 0.1
ep = altDemography.addEpoch(1000, newSize, 0.0)

for p in polyOp:
	print p



myPowers = []
myTimes = []
for e in ep:
	powerIs = myPower.getPower(altDemography, POWER_CALCULATION_SIZE)
	myPowers.append(powerIs)
	myTimes.append(e.myTime)
	print powerIs

for i in range(len(myTimes)):
	print '%1.4f \t%1.4f' % (myTimes[i], myPowers[i])

plt.semilogx( myTimes, myPowers )
plt.ylabel('Power')
plt.xlabel('Time of event (in coalescent time)')
plt.title('CC model with 10-fold growth. S=50, n=30')
plt.show()